package org.lep.leetcode.binarytreelevelordertraversal;

import org.lep.leetcode.binarytreeinordertraversal.BinaryTreeInOrderTraversal;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Stack;

/**
 * Source : https://oj.leetcode.com/problems/binary-tree-level-order-traversal/
 *
 * Created by lverpeng on 2017/8/12.
 *
 * Given a binary tree, return the level order traversal of its nodes' values.
 * (ie, from left to right, level by level).
 *
 * For example:
 * Given binary tree {3,9,20,#,#,15,7},
 *
 *     3
 *    / \
 *   9  20
 *     /  \
 *    15   7
 *
 * return its level order traversal as:
 *
 * [
 *   [3],
 *   [9,20],
 *   [15,7]
 * ]
 *
 */
public class BinaryTreeLevelOrderTraversal {

    /**
     * 按层次输出二叉树的遍历结果
     *
     * 使用BFS，借助栈实现，因为需要知道节点属于第几层，在入栈的时候就需要把节点和节点所在层level入栈
     *
     * @param root
     * @return
     */
    public List<List<Integer>> levelOrderTraversal (TreeNode root) {
        List<List<Integer>> result = new ArrayList<List<Integer>>();
        if (root == null) {
            return null;
        }
        Stack<LevelNode> stack = new Stack<LevelNode>();
        int currentLevel = 0;
        int lastLevel = 0;
        stack.push(new LevelNode(root, currentLevel));
        List<Integer> levelList = new ArrayList<Integer>();
        result.add(levelList);

        while (stack.size() > 0) {
            currentLevel = lastLevel;
            LevelNode node = stack.firstElement();
            stack.remove(0);
            if (node.node == null) {
                continue;
            }
            if (node.level != currentLevel) {
                levelList = new ArrayList<Integer>();
                result.add(levelList);
                lastLevel++;
            }
            result.get(lastLevel).add(node.node.value);

            stack.push(new LevelNode(node.node.leftChild, currentLevel + 1));
            stack.push(new LevelNode(node.node.rightChild, currentLevel + 1));
        }

        return result;
    }

    /**
     *
     * 按层次输出树
     *
     * 使用两个list分别保存当前层的node和下一层的node
     * 先将节点加入当前层
     * 遍历当前层，将每个节点的左右子节点加入下一层的list
     * 遍历完成之后，swap当前层和下一层的list，再次遍历
     *
     * @param root
     * @return
     */
    public List<List<Integer>> levelOrderTraversal1 (TreeNode root) {
        List<List<Integer>> result = new ArrayList<List<Integer>>();
        if (root == null) {
            return result;
        }
        List<TreeNode> curList = new ArrayList<TreeNode>();
        List<TreeNode> nextList = new ArrayList<TreeNode>();
        curList.add(root);

        while (curList.size() > 0) {
            List<Integer> list = new ArrayList<Integer>();
            for (int i = 0; i < curList.size(); i++) {
                if (curList.get(i) == null) {
                    continue;
                }
                list.add(curList.get(i).value);
                nextList.add(curList.get(i).leftChild);
                nextList.add(curList.get(i).rightChild);
            }
            if (list.size() > 0) {
                result.add(list);
            }
            curList.clear();
            curList.addAll(nextList);
            nextList.clear();
        }

        return result;

    }


    public static void print (List<List<Integer>> lists) {
        StringBuffer buffer = new StringBuffer();
        buffer.append("[\n");
        for (int i = 0; i < lists.size(); i++) {
            buffer.append(Arrays.toString(lists.get(i).toArray(new Integer[lists.get(i).size()])));
            buffer.append(",\n");
        }
        if (lists.size() > 0) {
            buffer.deleteCharAt(buffer.length()-2);
        }
        buffer.append("]");
        System.out.println(buffer.toString());
    }

    public TreeNode createTree (char[] treeArr) {
        TreeNode[] tree = new TreeNode[treeArr.length];
        for (int i = 0; i < treeArr.length; i++) {
            if (treeArr[i] == '#') {
                tree[i] = null;
                continue;
            }
            tree[i] = new TreeNode(treeArr[i]-'0');
        }
        int pos = 0;
        for (int i = 0; i < treeArr.length && pos < treeArr.length-1; i++) {
            if (tree[i] != null) {
                tree[i].leftChild = tree[++pos];
                if (pos < treeArr.length-1) {
                    tree[i].rightChild = tree[++pos];
                }
            }
        }
        return tree[0];
    }


    private class LevelNode {
        TreeNode node;
        int level;

        public LevelNode(TreeNode node, int level) {
            this.node = node;
            this.level = level;
        }
    }

    private class TreeNode {
        TreeNode leftChild;
        TreeNode rightChild;
        int value;

        public TreeNode(int value) {
            this.value = value;
        }

        public TreeNode() {
        }
    }

    public static void main(String[] args) {
        BinaryTreeLevelOrderTraversal binaryTreeLevelOrderTraversal = new BinaryTreeLevelOrderTraversal();
        char[] arr = new char[]{'3','9','2','#','#','1','7'};

        TreeNode tree = binaryTreeLevelOrderTraversal.createTree(arr);


        print(binaryTreeLevelOrderTraversal.levelOrderTraversal(tree));
        print(binaryTreeLevelOrderTraversal.levelOrderTraversal1(tree));

    }

}
